"""
Copyright (c) 2022 Guangdong University of Technology
PhotLab is licensed under [Open Source License].
You can use this software according to the terms and conditions of the [Open Source License].
You may obtain a copy of [Open Source License] at: [https://open.source.license/]

THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.

See the [Open Source License] for more details.

Author: Meng Xiang, Junjiang Xiang
Created: 2023/8/19
Supported by: National Key Research and Development Program of China
"""
import numpy as np

'''旧版 暂时保留 需修改
def cd_compensation(input,dispersion,Lambda,FiberLength,samplerate):
    """

    Args:
        input:
        dispersion: 色散 D
        Lambda: 1550nm
        FiberLength:光纤总长度
        samplerate: 采样率

    Returns:

    """
    c = 2.99792458e8
    D = dispersion * 1e-12 / (1e-9 * 1e3)
    Lambda = Lambda * 1e-9
    Beta2 = -D * Lambda ** 2 / (2 * np.pi * c)
    Fiber_Length = FiberLength * 1e3
    DTime = 1 / (samplerate)

    Number_of_Samples = len(input)
    MaxFreq = 0.5 / DTime
    DFreq = 2 * MaxFreq / (Number_of_Samples - 1)

    VFreq = np.arange(-1, Number_of_Samples-1) * DFreq
    VFreq = VFreq - 0.5 * max(VFreq)
    VOmeg = 2 * np.pi * VFreq

    Disper = (1j / 2) * Beta2 * VOmeg ** 2 * Fiber_Length

    Freq_Samples = np.fft.fftshift(np.fft.fft(input.T))
    Output_Freq_Samples = Freq_Samples * np.exp(Disper)
    output = np.fft.ifft(np.fft.ifftshift(Output_Freq_Samples))

    return output.T
'''

def Getft(Signal):
    if Signal["FFT_len"]%2 == 0:    # 求余
        fft_n = Signal["FFT_len"]/2
        fft_n = int(fft_n)
        v1 = np.linspace(start=-fft_n,stop=fft_n,num=2*fft_n,dtype=float)
        f_re = (Signal["Sampling_rate"]*v1)/Signal["FFT_len"]
    else:
        fft_n = (Signal["FFT_len"]-1)/2
        fft_n = int(fft_n)
        # v2 = np.array([range(0-fft_n, fft_n+1)], dtype=float)
        v2 = np.linspace(start=-fft_n,stop=fft_n,num=2*fft_n,dtype=float)
        f_re = (Signal["Sampling_rate"]*v2)/Signal["FFT_len"]
    ft = np.reshape(f_re, len(f_re))
    return ft

def cd_compensation(input, dispersion, Lambda, FiberLength, sampling_rate, num_symbols, up_sampling_factor, beta2):
    """

    Args:
        input:
        dispersion: 色散 D
        Lambda: 1550nm
        FiberLength:光纤总长度
        samplerate: 采样率

    Returns:

    """

    c = 2.99792458e8
    D = dispersion * 1e-12 / (1e-9)
    Lambda = Lambda * 1e-9
    Beta2 = -D * Lambda ** 2 / (2 * np.pi * c)
    # Fiber_Length = FiberLength * 1e3
    Fiber_Length = FiberLength
    DTime = 1 / (sampling_rate)


    Number_of_Samples = len(input[0])
    MaxFreq = 0.5 / DTime
    DFreq = 2 * MaxFreq / (Number_of_Samples - 1)

    VFreq = np.arange(-1, Number_of_Samples-1) * DFreq
    VFreq = VFreq - 0.5 * max(VFreq)
    VOmeg = 2 * np.pi * VFreq

    Signal_Para = {}
    Signal_Para["FFT_len"] = num_symbols * up_sampling_factor
    Signal_Para["Sampling_rate"] = sampling_rate
    ft = Getft(Signal_Para)
    f_omega = 2 * np.pi * ft
    print(f'f_omega  {f_omega}')
    # Disper = (1j / 2) * Beta2 * VOmeg ** 2 * Fiber_Length
    Disper = (1j / 2) * beta2 * f_omega ** 2 * Fiber_Length

    data_x = input[0].reshape(-1)
    data_y = input[1].reshape(-1)

    # Freq_Samples_X = np.fft.fftshift(np.fft.fft(input[0].T))
    # Output_Freq_Samples_X = Freq_Samples_X * np.exp(Disper)
    # output_X = np.fft.ifft(np.fft.ifftshift(Output_Freq_Samples_X))
    #
    # Freq_Samples_Y = np.fft.fftshift(np.fft.fft(input[1].T))
    # Output_Freq_Samples_Y = Freq_Samples_Y * np.exp(Disper)
    # output_Y = np.fft.ifft(np.fft.ifftshift(Output_Freq_Samples_Y))

    Freq_Samples_X = np.fft.fftshift(np.fft.fft(data_x))
    Output_Freq_Samples_X = Freq_Samples_X * np.exp(Disper)
    output_X = np.fft.ifft(np.fft.ifftshift(Output_Freq_Samples_X))

    Freq_Samples_Y = np.fft.fftshift(np.fft.fft(data_y))
    Output_Freq_Samples_Y = Freq_Samples_Y * np.exp(Disper)
    output_Y = np.fft.ifft(np.fft.ifftshift(Output_Freq_Samples_Y))

    output_X = np.expand_dims(output_X, axis=1)
    output_Y = np.expand_dims(output_Y, axis=1)

    return [output_X, output_Y]